# -*- coding:utf-8 -*-
"""
# @file name        : lesson-06-1_split_dataset
# @author           : QuZhang
# @data             : 2020-11-28 22:44
# @brief            : 将数据集划分为训练集，验证集，测试集
"""

import random
import shutil
import os
abspath = os.path.abspath(__file__)
# print("当前文件的完整路径：", abspath)
BASE_DIR = os.path.dirname(abspath)  # 定位到当前目录,abspath的上一级路径
# print("当前文件夹路径：", BASE_DIR)  # D:\CV_AMBITION\pytorch\deepshare\PyTorch\code\lesson\lesson-06

def makedir(new_dir):
    if not os.path.exists(new_dir):
        os.makedirs(new_dir)

if __name__ == '__main__':

    # 一、构建存储数据的文件夹
    # 1.获取原始数据所在文件夹路径
    dataset_dir = os.path.abspath(os.path.join(BASE_DIR, "..", "..", "data", "RMB_data"))
    # 2.获取划分后数据所在文件夹路径
    split_dir = os.path.abspath(os.path.join(BASE_DIR, '..', "..", "data", 'rmb_split'))
    train_dir = os.path.join(split_dir, "train")  # 训练集数据
    valid_dir = os.path.join(split_dir, 'valid')  # 验证集路径
    test_dir = os.path.join(split_dir, "test")  # 测试集路径

    if not os.path.exists(dataset_dir):
        raise Exception("\n{} 不存在，请下载 02-01-数据-RMB_data.rar 放到\n{} 下,并解压即可".format(
            dataset_dir, os.path.dirname(dataset_dir)
        ))

    # 二、设置数据集划分比例
    train_pct = 0.8
    valid_pct = 0.1
    test_pct = 0.1

    # 深度优先搜索遍历某个路径下的所有目录和文件
    for root, dirs, files in os.walk(dataset_dir):
        """
        root: 当前目录的绝对路径
        dirs: root里有哪些目录(子目录)
        files: root里有哪些文件
        """
        # 逐个遍历所有的子目录
        for sub_dir in dirs:
            sub_dir_abspath = os.path.join(root, sub_dir)
            print("sub_dir:", sub_dir)
            imgs = os.listdir(sub_dir_abspath)  # root/sub_dir目录下的所有文件列表
            imgs = list(filter(lambda x: x.endswith('.jpg'), imgs))  # 过滤掉非.jpg后缀的图片
            random.shuffle(imgs)  # 随机打乱图片顺序
            img_count = len(imgs)

            # 确定每种数据集的图片张数
            train_point = int(img_count * train_pct)  # 设置划分的停止点
            valid_point = int(img_count * (train_pct + valid_pct))
            # print("train size: ", train_point, end=' ')
            # print("valid size: ", valid_point-train_point, end=' ')
            # print("test size: ", img_count-valid_point)

            if img_count == 0:
                print("{}目录下，无图片，请检查".format(sub_dir_abspath))
                import sys
                sys.exit(0)

            for i in range(img_count):
                if i < train_point:
                    # 训练集
                    out_dir = os.path.join(train_dir, sub_dir)
                elif i < valid_point:
                    # 验证集
                    out_dir = os.path.join(valid_dir, sub_dir)
                else:
                    out_dir = os.path.join(test_dir, sub_dir)

                makedir(out_dir)  # 创建目录

                target_path = os.path.join(out_dir, imgs[i])  # 完整的图片路径
                src_path = os.path.join(dataset_dir, sub_dir, imgs[i])

                shutil.copy(src_path, target_path)  # 拷贝文件的内容和权限

            print('Class:{}, train:{}, \nvalid:{}, \ntest:{}'.format(sub_dir, train_dir, valid_dir, test_dir))
            print("已在 {} 创建划分好的数据".format(out_dir))